import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

def weibull_cdf(t, eta, beta):
    return 1.0 - torch.exp(-1.0 * (t / eta) ** beta)

class WeiBullCDF(nn.Module):
    
    def __init__(self, eta, beta):
        super().__init__()
        self.eta = eta
        self.beta = beta
    
    def forward(self, inputs):
        return 1.0 - torch.exp(-1.0 * (inputs / self.eta) ** self.beta)


class WeiBullLossRMSE(nn.Module):
    
    def __init__(self, full_life, lambda_mod, eta, beta, eps=1e-8):
        super().__init__()
        self.full_life = full_life
        self.lambda_mod = lambda_mod
        self.weibull_cdf = WeiBullCDF(eta, beta)
        self.eps = eps
        
    def forward(self, inputs, target):
        target_age = self.full_life - target
        input_age = self.full_life - inputs
        input_age = input_age[torch.isfinite(input_age)]

        cdf_true = self.weibull_cdf(target_age)
        cdf_pred = self.weibull_cdf(input_age)
        
        return torch.sqrt(torch.mean(inputs - target)**2 + self.eps) + self.lambda_mod * torch.sqrt(torch.mean(cdf_pred - cdf_true)**2 + self.eps)
        
    
    
class WeiBullLossMSE(nn.Module):
    
    def __init__(self, full_life, lambda_mod, eta, beta):
        super().__init__()
        self.full_life = full_life
        self.lambda_mod = lambda_mod
        self.weibull_cdf = WeiBullCDF(eta, beta)
        
    def forward(self, inputs, target):
        target_age = self.full_life - target
        input_age = self.full_life - inputs
        input_age = input_age[torch.isfinite(input_age)]
        
        cdf_true = self.weibull_cdf(target_age)
        cdf_pred = self.weibull_cdf(input_age)
        
        #return self.lambda_mod * torch.mean((cdf_pred - cdf_true) ** 2)
        return torch.mean(torch.square(inputs - target)) + self.lambda_mod * torch.mean(torch.square(cdf_pred - cdf_true))
        
        
class WeiBullLossRMSLE(nn.Module):
    
    def __init__(self, full_life, lambda_mod, eta, beta, eps=1e-8):
        super().__init__()
        self.full_life = full_life
        self.lambda_mod=lambda_mod
        self.weibull_cdf = WeiBullCDF(eta, beta)
        self.eps = eps
    
    
    def forward(self, input, target):
        target_age = self.full_life - target
        input_age = self.full_life - input
        input_age = input_age[torch.isfinite(input_age)]
        
        cdf_true = self.weibull_cdf(target_age)
        cdf_pred = self.weibull_cdf(input_age)
        
        return self.lambda_mod * torch.sqrt(torch.mean(torch.log(cdf_pred+1) -  torch.log(cdf_true+1)) ** 2 + self.eps)
    
class RMSELoss(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, inputs, targets):
        return torch.sqrt(torch.mean(torch.square(inputs - targets)))

class MAELoss(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, inputs, targets):
        return torch.mean(torch.abs(inputs - targets))

class ChangedMse(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, inputs, targets):
        sum = np.zeros_like(inputs.cpu().detach().numpy())
        sum = torch.tensor(sum, dtype=torch.float32).to(inputs.device)
        for i in range(inputs.size(0)):
            if inputs[i] > targets[i]:
                sum[i] = 1.3 * torch.square(inputs[i] - targets[i])
            else:
                sum[i] = torch.square(targets[i] - inputs[i])

        return sum.mean()
    
    
def get_loss_fn(loss_kind, full_life=1, lambda_mod=0.2, eta=1.4, beta=2):
    loss_kind = loss_kind.lower()
    if loss_kind == "weibull_mse":
        loss_fn = WeiBullLossMSE(full_life, lambda_mod, eta, beta)
    elif loss_kind == "weibull_rmse":
        loss_fn = WeiBullLossRMSE(full_life, lambda_mod, eta, beta)
    elif loss_kind == "weibull_rmsle":
        loss_fn = WeiBullLossRMSLE(full_life, lambda_mod, eta, beta)
    elif loss_kind == "mse":
        loss_fn = nn.MSELoss()
    elif loss_kind == "c_mse":
        loss_fn = ChangedMse()
    elif loss_kind == "rmse":
        loss_fn = RMSELoss()
    elif loss_kind == "mae":
        loss_fn = MAELoss()
    else:
        raise ValueError(f"no such loss function called {loss_kind}.")
    

    return loss_fn

